import os
import copy
import json
import math
import argparse
import numpy as np
import matplotlib.pyplot as plt
import random

def calc_plane(bgn, end, step):
    x = np.arange(bgn, end, step)
    y = np.arange(bgn, end, step)
    X, Y = np.meshgrid(x, y)

    #a, b, c, d = 1, 10, 2, 3
    if isinstance(X, np.ndarray):
        z = []
        for i in range(X.shape[0]):
            for j in range(X.shape[1]):
                tx, ty = X[i][j], Y[i][j]
                #tz = (1-ty**5+tx**5)*np.exp(-tx**2-ty**2)
                #tz = -np.cos(tx)*np.cos(ty)-np.cos(np.pi-tx**3)+np.sin(ty**5)-np.cos(ty**np.pi)
                tz = tx + ty + np.sin(tx) * tx/2 + np.sin(ty) * ty/2
                #tz = tx + ty + random.randint(1, 10)
                z.append(tz)
                #tz.append(random.randint(tx, ty + 100))
                #tz.append(-1 * (a * X[i][j] + b * Y[i][j] + d) / c)
        Z = np.array(z).reshape(X.shape)
    return X, Y, Z

if __name__=='__main__':
    #bgn, end = 1.1, 3.4
    bgn, end = -10, 10.1

    # 生成组织曲面
    #X1, Y1, Z1 = calc_plane(bgn, end, 0.05)
    X1, Y1, Z1 = calc_plane(bgn, end, 0.01)
    print('shape1: ({}, {}, {})'.format(X1.shape, Y1.shape, Z1.shape))

    # 生成对焦平面
    #X2, Y2, Z2 = calc_plane(bgn, end, 0.285)
    X2, Y2, Z2 = calc_plane(bgn, end, 2.5)
    print('shape2: ({}, {}, {})'.format(X2.shape, Y2.shape, Z2.shape))

    # 生成对焦点
    X3, Y3, Z3, C3 = [], [], [], []
    for i in range(X2.shape[0]):
        for j in range(X2.shape[1]):
            tx, ty, tz = X2[i][j], Y2[i][j], Z2[i][j]
            X3.append(tx)
            Y3.append(ty)
            Z3.append(tz)
            C3.append('r')

    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')
    #ax.set_xlabel('X')
    #ax.set_ylabel('Y')
    #ax.set_zlabel('Z')

    ax.plot_surface(X1, Y1, Z1, color='w', alpha=0.5)
    ax.plot_surface(X2, Y2, Z2, color='r', alpha=0.3)
    ax.scatter(X3, Y3, Z3, c=C3)

    plt.show()
